import {
  CreateImageResult,
  GetImageGenerationLocationResult,
  UpLoadImageResult,
} from '@/api/comfyUI';
import { MessageType } from '@/config';
import { createWorkFlowApi } from '@/lib/comfyUI';
import { generateRandomId, getDataUrl, getRandomElement, joinLink } from '@/lib/utils';
import { BackgroundReference, useSettingsStore } from '@/store';

export function useComfyUI() {
  const processSettings = useSettingsStore((state) => state.processSettings);

  const upLoadImage = useCallback(
    async (blob: Blob) => {
      const imageDataUrl = await getDataUrl(blob);
      const result = await browser.runtime.sendMessage({
        type: MessageType.corssRequest,
        data: {
          kind: 'comfyUiUploadImage',
          config: {
            url: joinLink(processSettings.comfyUILink, 'api/upload/image'),
            dataUrl: imageDataUrl,
            name: `${generateRandomId()}.jpg`,
          },
        },
      });
      if (result.code === 0) {
        const data = result.data as UpLoadImageResult;
        return data.name;
      } else {
        throw new Error(result.message);
      }
    },
    [processSettings],
  );

  const queryImageUrl = useCallback(
    async (promptId: string) => {
      return new Promise((resolve, reject) => {
        const timer = setInterval(async () => {
          const result = await browser.runtime.sendMessage({
            type: MessageType.corssRequest,
            data: {
              kind: 'comfyUiGetImageGenerationLocation',
              config: {
                url: joinLink(processSettings.comfyUILink, `history/${promptId}`),
              },
            },
          });
          if (result.code === 0) {
            const data: GetImageGenerationLocationResult = result.data;
            let url = '';
            try {
              url =
                data[promptId].outputs[processSettings.comfyUIOutputImageNodeId].images?.[0]
                  .filename || '';
              if (url) {
                clearInterval(timer);
                resolve(url);
              }
            } catch (error) {
              // console.error(error);
            }
          } else {
            clearInterval(timer);
            reject(result.message);
          }
          // console.log(result);
        }, 3000);
        setTimeout(() => {
          reject(new Error('获取替换背景的图片超时'));
          clearInterval(timer);
        }, 1000 * processSettings.aiProcessTimeout);
      });
    },
    [processSettings],
  );

  const generateAIImage = useCallback(
    async (
      sourceImage: Blob,
    ): Promise<{
      resultUrl: string;
      resultBlob: Blob;
    }> => {
      // 第一步上传图像和背景
      const backgroundReference = getRandomElement<BackgroundReference>(
        processSettings.backgroundReferenceList,
      );
      if (!backgroundReference) {
        throw new Error('至少需要设置一个背景参考');
      }
      const backgroundBlob = await (await fetch(backgroundReference.dataUrl)).blob();
      const [inputImage, backgroundImage] = await Promise.all([
        upLoadImage(sourceImage),
        upLoadImage(backgroundBlob),
      ]);

      // console.log(inputImage, backgroundImage);
      // 创建对象
      const newWorkFlowApi = createWorkFlowApi({
        inputImage,
        backgroundImage,
        comfyUIInputImageNodeId: processSettings.comfyUIInputImageNodeId,
        comfyUIBackgroundImageNodeId: processSettings.comfyUIInputBackgroundImageNodeId,
      });
      const result = await browser.runtime.sendMessage({
        type: MessageType.corssRequest,
        data: {
          kind: 'comfyUiCreateImage',
          config: {
            url: joinLink(processSettings.comfyUILink, 'prompt'),
            prompt: newWorkFlowApi,
          },
        },
      });
      if (result.code === 0) {
        const data: CreateImageResult = result.data;
        const promptId = data.prompt_id;
        // 轮询获取最终的生成图片位置
        const url = await queryImageUrl(promptId);
        const viewUrl = joinLink(processSettings.comfyUILink, `view?filename=${url}`);
        const downloadImageResult = await browser.runtime.sendMessage({
          type: MessageType.corssRequest,
          data: {
            kind: 'comfyUiDownloadImage',
            config: {
              url: viewUrl,
            },
          },
        });
        if (downloadImageResult.code === 0) {
          const dataUrl: string = downloadImageResult.data;
          const imageBlob = await (await fetch(dataUrl)).blob();
          return {
            resultUrl: viewUrl,
            resultBlob: imageBlob,
          };
        } else {
          throw new Error(downloadImageResult.message);
        }
      } else {
        throw new Error(result.message);
      }
      // console.log(result);
    },
    [upLoadImage, processSettings, queryImageUrl],
  );

  const generateAIImageLoop = useCallback(
    async (
      sourceImage: Blob,
      aiRetryTime: number,
      retryCallback: (error: unknown) => void,
    ): Promise<{
      resultUrl: string;
      resultBlob: Blob;
    }> => {
      try {
        return await generateAIImage(sourceImage);
      } catch (error) {
        console.log(error);
        if (aiRetryTime > 0) {
          retryCallback(error);
          return generateAIImageLoop(sourceImage, aiRetryTime - 1, retryCallback);
        } else {
          throw new Error(`请确保 ComfyUI 服务正在运行: ${(error as Error).message}`);
        }
      }
    },
    [generateAIImage],
  );

  return {
    generateAIImage,
    generateAIImageLoop,
  };
}
